# Terrell 旋转模拟器！
# 画质略低于80年代红白机
# 参考：https://www.bilibili.com/video/BV1JY411L7xm
# 他的实现应该好看+快得多
# Gitee Repo

import numpy as np
import multiprocessing
import multiprocessing.shared_memory
import time
import sys

def generate_cube_data(center, length, width, height):
    """
    生成长方体的顶点和边数据
    
    参数:
    center: 长方体中心坐标 [x, y, z]
    length: 长度 (x轴方向)
    width: 宽度 (y轴方向)  
    height: 高度 (z轴方向)
    
    返回:
    vertices: 8个顶点坐标 (8×3)
    edge_start: 边起点坐标 (12×3)
    edge_end: 边终点坐标 (12×3)
    """
    cx, cy, cz = center
    l, w, h = length/2, width/2, height/2
    
    # 生成8个顶点
    vertices = np.array([
        [cx - l, cy - w, cz - h],  # 0: 左下后
        [cx + l, cy - w, cz - h],  # 1: 右下后
        [cx + l, cy + w, cz - h],  # 2: 右上后
        [cx - l, cy + w, cz - h],  # 3: 左上后
        [cx - l, cy - w, cz + h],  # 4: 左下前
        [cx + l, cy - w, cz + h],  # 5: 右下前
        [cx + l, cy + w, cz + h],  # 6: 右上前
        [cx - l, cy + w, cz + h]   # 7: 左上前
    ])
    
    # 定义12条边的连接关系
    edges = [
        (0, 1), (1, 2), (2, 3), (3, 0),  # 底面边
        (4, 5), (5, 6), (6, 7), (7, 4),  # 顶面边  
        (0, 4), (1, 5), (2, 6), (3, 7)   # 垂直边
    ]
    
    # 生成边的起点和终点坐标
    edge_start = np.array([vertices[i] for i, j in edges])
    edge_end = np.array([vertices[j] for i, j in edges])
    
    return vertices, edge_start, edge_end

def points_near_segment(points, segment_start, segment_end, threshold):
    """
    批量判断多个点是否在线段附近
    
    参数:
    points: numpy数组, 点的坐标数组, 形状为 (n, 3)
    segment_start: numpy数组, 线段起点 [x, y, z]
    segment_end: numpy数组, 线段终点 [x, y, z]
    threshold: float, 距离阈值
    
    返回:
    numpy数组: 布尔数组，表示每个点是否在线段附近
    """
    # 将输入转换为numpy数组
    if points.ndim != 2 or points.shape[1] != 3 or segment_start.ndim != 1 or segment_start.shape[0] != 3 or segment_end.ndim != 1 or segment_end.shape[0] != 3:
        raise ValueError('Wrong Shape')
    
    P = points
    A = segment_start[None,:]
    B = segment_end[None,:]
    
    # 计算向量
    AB = B - A  # 线段方向向量
    AP = P - A  # 每个点到起点的向量
    
    # 计算投影参数t
    t = np.dot(AP, AB.T) / np.dot(AB, AB.T)
    #print(AP.shape)
    #print(AB.shape)
    #print(t.shape)

    # 限制t在[0,1]范围内
    t = np.clip(t, 0, 1)
    
    # 计算线段上距离每个点最近的点
    closest_points = A + t * AB
    #print(A.shape)
    #print(t.shape)
    #print(AB.shape)
    #print(closest_points.shape)

    # 计算每个点到最近点的距离
    distances = np.linalg.norm(P - closest_points, axis=1)
    
    # 判断距离是否在阈值内
    return distances <= threshold

def points_near_any_segment(points, edge_start, edge_end, threshold):
    res = None
    for i in range(edge_start.shape[0]):
        _res = points_near_segment(points, edge_start[i,:], edge_end[i,:], threshold)
        if res is None:
            res = _res.copy()
        else:
            res = res | _res
    return res
    
def pcolor(x,y,z):
    import numpy as np
    import matplotlib.pyplot as plt

    # 创建图形
    plt.figure(figsize=(8, 6))
    # 使用pcolormesh绘制伪彩色图（推荐使用，比pcolor更高效）
    pc = plt.pcolormesh(x, y, z, shading='auto', cmap='viridis')
    plt.colorbar(pc, label='Value')  # 添加颜色条
    # 设置标题和坐标轴标签
    plt.title('Python pcolormesh (similar to MATLAB pcolor)')
    plt.xlabel('X axis')
    plt.ylabel('Y axis')
    # 显示图形
    plt.show()

vertices, edge_start, edge_end = generate_cube_data(np.array([0,0,-2]), 2, 2, 2)

L = 1;
dx = 0.005;
dt = 0.025

b = 0.8
g = 1/np.sqrt(1-b**2)
c = 1;

[x1, y1] = np.meshgrid(np.arange(-L,L,dx),np.arange(-L,L,dx))
n = x1.shape[0]

args = sys.argv
t = float(args[1]) if len(args) > 1 else 1
X1_base = np.vstack((c*t*np.ones((1,n**2)),x1.flatten(),y1.flatten(),np.zeros((1,n**2))))

#dX = np.array([[-c],[0],[0],[-c]])
dX = X1_base[1:4,:] - np.array([[0],[0],[0.5]])
dX = dX / np.linalg.norm(dX,axis=0)
dX = np.vstack((-c*np.ones((1,n**2)),c*dX))
print(dX)

Lorentz = np.array([[g,-g*b,0,0],[-g*b,g,0,0],[0,0,1,0],[0,0,0,1]])

def worker(i,n_proc,_tick_start,_tick_end,shm_name):
    shm = multiprocessing.shared_memory.SharedMemory(name=shm_name)
    shared_res = np.ndarray((X1_base.shape[1],n_proc), dtype=np.bool, buffer=shm.buf)
    
    res = None
    for _tick in range(_tick_start,_tick_end):
        _t = _tick*dt
        _X1 = X1_base + dX * _t
        _X2 = Lorentz @ _X1
        
        _X2_vec = _X2[1:4,:].T
        _res = points_near_any_segment(_X2_vec, edge_start, edge_end, 0.1)
        
        if res is None:
            res = _res.copy()
        else:
            res = res | _res

        print(_tick)
    
    shared_res[:,i] = res.astype(np.bool).copy()
    shm.close()
    return


if __name__ == '__main__':
    n_proc = 8
    shm = multiprocessing.shared_memory.SharedMemory(create=True, size=X1_base.shape[1]*n_proc)
    shared_res = np.ndarray((X1_base.shape[1],n_proc), dtype=np.bool, buffer=shm.buf)

    _TICK = 250
    step = _TICK // n_proc

    start_time = time.time()

    process = []
    for i in range(n_proc):
        p = multiprocessing.Process(target=worker,args=(i,n_proc,i*step,(i+1)*step,shm.name))
        process.append(p)
        p.start()
        
    
    for p in process:
        p.join()
    
    res = shared_res.copy()
    res = np.bitwise_or.reduce(res, axis=1)

    shm.close()
    shm.unlink()

    end_time = time.time()
    print(end_time-start_time)

    fig = res.reshape(n,n)
    #pcolor(x1,y1,fig)

    import PIL
    import PIL.Image
    image = PIL.Image.fromarray(fig)
    image.save(f'terrell_{t:0.2f}.png')

